{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Display matplotlib animations as HTML5 video\n",
    "\n",
    "Based on the source in https://github.com/jakevdp/JSAnimation/blob/master/JSAnimation/IPython_display.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import random\n",
    "import string\n",
    "import warnings\n",
    "if sys.version_info < (3, 0):\n",
    "    from cStringIO import StringIO as InMemory\n",
    "else:\n",
    "    from io import BytesIO as InMemory\n",
    "from matplotlib.animation import writers, FileMovieWriter\n",
    "from base64 import b64encode"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Reusable animation objects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def circle(ax, x, y, r, color='r'):\n",
    "    count = 20\n",
    "    return ax.fill([x + r * np.cos(2.* i * np.pi/count) for i in range(count+1)],\n",
    "            [y + r * np.sin(2.* i * np.pi/count) for i in range(count+1)],\n",
    "            color)\n",
    "\n",
    "def line(ax, x1, y1, x2, y2, color='black'):\n",
    "    norm = np.sqrt((y2-y1)**2 + (x2-x1)**2)\n",
    "    normal = [(y2-y1)/norm, -(x2-x1)/norm]\n",
    "    return ax.fill(\n",
    "    [x1-0.01*normal[0],x1+0.01*normal[0],x2+0.01*normal[0],x2-0.01*normal[0]],\n",
    "    [y1-0.01*normal[1],y1+0.01*normal[1],y2+0.01*normal[1],y2-0.01*normal[1]],\n",
    "            color)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    ICON_DIR = os.path.join(os.path.dirname(__file__), 'icons')\n",
    "except NameError:\n",
    "    ICON_DIR = ''\n",
    "\n",
    "class _Icons(object):\n",
    "    \"\"\"This class is a container for base64 representations of the icons\"\"\"\n",
    "    icons = ['first', 'prev', 'reverse', 'pause', 'play', 'next', 'last']\n",
    "\n",
    "    def __init__(self, icon_dir=ICON_DIR, extension='png'):\n",
    "        self.icon_dir = icon_dir\n",
    "        self.extension = extension\n",
    "        for icon in self.icons:\n",
    "            setattr(self, icon,\n",
    "                    self._load_base64('{0}.{1}'.format(icon, extension)))\n",
    "\n",
    "    def _load_base64(self, filename):\n",
    "        data = open(os.path.join(self.icon_dir, filename), 'rb').read()\n",
    "        return 'data:image/{0};base64,{1}'.format(self.extension,\n",
    "                                                  b64encode(data).decode('ascii'))\n",
    "\n",
    "\n",
    "JS_INCLUDE = \"\"\"\n",
    "<script language=\"javascript\">\n",
    "  /* Define the Animation class */\n",
    "  function Animation(frames, img_id, slider_id, interval, loop_select_id){\n",
    "    this.img_id = img_id;\n",
    "    this.slider_id = slider_id;\n",
    "    this.loop_select_id = loop_select_id;\n",
    "    this.interval = interval;\n",
    "    this.current_frame = 0;\n",
    "    this.direction = 0;\n",
    "    this.timer = null;\n",
    "    this.frames = new Array(frames.length);\n",
    "    for (var i=0; i<frames.length; i++)\n",
    "    {\n",
    "     this.frames[i] = new Image();\n",
    "     this.frames[i].src = frames[i];\n",
    "    }\n",
    "    document.getElementById(this.slider_id).max = this.frames.length - 1;\n",
    "    this.set_frame(this.current_frame);\n",
    "  }\n",
    "  Animation.prototype.get_loop_state = function(){\n",
    "    var button_group = document[this.loop_select_id].state;\n",
    "    for (var i = 0; i < button_group.length; i++) {\n",
    "        var button = button_group[i];\n",
    "        if (button.checked) {\n",
    "            return button.value;\n",
    "        }\n",
    "    }\n",
    "    return undefined;\n",
    "  }\n",
    "  Animation.prototype.set_frame = function(frame){\n",
    "    this.current_frame = frame;\n",
    "    document.getElementById(this.img_id).src = this.frames[this.current_frame].src;\n",
    "    document.getElementById(this.slider_id).value = this.current_frame;\n",
    "  }\n",
    "  Animation.prototype.next_frame = function()\n",
    "  {\n",
    "    this.set_frame(Math.min(this.frames.length - 1, this.current_frame + 1));\n",
    "  }\n",
    "  Animation.prototype.previous_frame = function()\n",
    "  {\n",
    "    this.set_frame(Math.max(0, this.current_frame - 1));\n",
    "  }\n",
    "  Animation.prototype.first_frame = function()\n",
    "  {\n",
    "    this.set_frame(0);\n",
    "  }\n",
    "  Animation.prototype.last_frame = function()\n",
    "  {\n",
    "    this.set_frame(this.frames.length - 1);\n",
    "  }\n",
    "  Animation.prototype.slower = function()\n",
    "  {\n",
    "    this.interval /= 0.7;\n",
    "    if(this.direction > 0){this.play_animation();}\n",
    "    else if(this.direction < 0){this.reverse_animation();}\n",
    "  }\n",
    "  Animation.prototype.faster = function()\n",
    "  {\n",
    "    this.interval *= 0.7;\n",
    "    if(this.direction > 0){this.play_animation();}\n",
    "    else if(this.direction < 0){this.reverse_animation();}\n",
    "  }\n",
    "  Animation.prototype.anim_step_forward = function()\n",
    "  {\n",
    "    this.current_frame += 1;\n",
    "    if(this.current_frame < this.frames.length){\n",
    "      this.set_frame(this.current_frame);\n",
    "    }else{\n",
    "      var loop_state = this.get_loop_state();\n",
    "      if(loop_state == \"loop\"){\n",
    "        this.first_frame();\n",
    "      }else if(loop_state == \"reflect\"){\n",
    "        this.last_frame();\n",
    "        this.reverse_animation();\n",
    "      }else{\n",
    "        this.pause_animation();\n",
    "        this.last_frame();\n",
    "      }\n",
    "    }\n",
    "  }\n",
    "  Animation.prototype.anim_step_reverse = function()\n",
    "  {\n",
    "    this.current_frame -= 1;\n",
    "    if(this.current_frame >= 0){\n",
    "      this.set_frame(this.current_frame);\n",
    "    }else{\n",
    "      var loop_state = this.get_loop_state();\n",
    "      if(loop_state == \"loop\"){\n",
    "        this.last_frame();\n",
    "      }else if(loop_state == \"reflect\"){\n",
    "        this.first_frame();\n",
    "        this.play_animation();\n",
    "      }else{\n",
    "        this.pause_animation();\n",
    "        this.first_frame();\n",
    "      }\n",
    "    }\n",
    "  }\n",
    "  Animation.prototype.pause_animation = function()\n",
    "  {\n",
    "    this.direction = 0;\n",
    "    if (this.timer){\n",
    "      clearInterval(this.timer);\n",
    "      this.timer = null;\n",
    "    }\n",
    "  }\n",
    "  Animation.prototype.play_animation = function()\n",
    "  {\n",
    "    this.pause_animation();\n",
    "    this.direction = 1;\n",
    "    var t = this;\n",
    "    if (!this.timer) this.timer = setInterval(function(){t.anim_step_forward();}, this.interval);\n",
    "  }\n",
    "  Animation.prototype.reverse_animation = function()\n",
    "  {\n",
    "    this.pause_animation();\n",
    "    this.direction = -1;\n",
    "    var t = this;\n",
    "    if (!this.timer) this.timer = setInterval(function(){t.anim_step_reverse();}, this.interval);\n",
    "  }\n",
    "</script>\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "DISPLAY_TEMPLATE = \"\"\"\n",
    "<div class=\"animation\" align=\"center\">\n",
    "    <img id=\"_anim_img{id}\">\n",
    "    <br>\n",
    "    <input id=\"_anim_slider{id}\" type=\"range\" style=\"width:350px\" name=\"points\" min=\"0\" max=\"1\" step=\"1\" value=\"0\" onchange=\"anim{id}.set_frame(parseInt(this.value));\"></input>\n",
    "    <br>\n",
    "    <button onclick=\"anim{id}.slower()\">&#8211;</button>\n",
    "    <button onclick=\"anim{id}.first_frame()\"><img class=\"anim_icon\" src=\"{icons.first}\"></button>\n",
    "    <button onclick=\"anim{id}.previous_frame()\"><img class=\"anim_icon\" src=\"{icons.prev}\"></button>\n",
    "    <button onclick=\"anim{id}.reverse_animation()\"><img class=\"anim_icon\" src=\"{icons.reverse}\"></button>\n",
    "    <button onclick=\"anim{id}.pause_animation()\"><img class=\"anim_icon\" src=\"{icons.pause}\"></button>\n",
    "    <button onclick=\"anim{id}.play_animation()\"><img class=\"anim_icon\" src=\"{icons.play}\"></button>\n",
    "    <button onclick=\"anim{id}.next_frame()\"><img class=\"anim_icon\" src=\"{icons.next}\"></button>\n",
    "    <button onclick=\"anim{id}.last_frame()\"><img class=\"anim_icon\" src=\"{icons.last}\"></button>\n",
    "    <button onclick=\"anim{id}.faster()\">+</button>\n",
    "  <form action=\"#n\" name=\"_anim_loop_select{id}\" class=\"anim_control\">\n",
    "    <input type=\"radio\" name=\"state\" value=\"once\" {once_checked}> Once </input>\n",
    "    <input type=\"radio\" name=\"state\" value=\"loop\" {loop_checked}> Loop </input>\n",
    "    <input type=\"radio\" name=\"state\" value=\"reflect\" {reflect_checked}> Reflect </input>\n",
    "  </form>\n",
    "</div>\n",
    "<script language=\"javascript\">\n",
    "  /* Instantiate the Animation class. */\n",
    "  /* The IDs given should match those used in the template above. */\n",
    "  (function() {{\n",
    "    var img_id = \"_anim_img{id}\";\n",
    "    var slider_id = \"_anim_slider{id}\";\n",
    "    var loop_select_id = \"_anim_loop_select{id}\";\n",
    "    var frames = new Array({Nframes});\n",
    "    {fill_frames}\n",
    "    /* set a timeout to make sure all the above elements are created before\n",
    "       the object is initialized. */\n",
    "    setTimeout(function() {{\n",
    "        anim{id} = new Animation(frames, img_id, slider_id, {interval}, loop_select_id);\n",
    "    }}, 0);\n",
    "  }})()\n",
    "</script>\n",
    "\"\"\"\n",
    "\n",
    "INCLUDED_FRAMES = \"\"\"\n",
    "  for (var i=0; i<{Nframes}; i++){{\n",
    "    frames[i] = \"{frame_dir}/frame\" + (\"0000000\" + i).slice(-7) + \".{frame_format}\";\n",
    "  }}\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "def _included_frames(frame_list, frame_format):\n",
    "    \"\"\"frame_list should be a list of filenames\"\"\"\n",
    "    return INCLUDED_FRAMES.format(Nframes=len(frame_list),\n",
    "                                  frame_dir=os.path.dirname(frame_list[0]),\n",
    "                                  frame_format=frame_format)\n",
    "\n",
    "\n",
    "def _embedded_frames(frame_list, frame_format):\n",
    "    \"\"\"frame_list should be a list of base64-encoded png files\"\"\"\n",
    "    template = '  frames[{0}] = \"data:image/{1};base64,{2}\"\\n'\n",
    "    embedded = \"\\n\"\n",
    "    for i, frame_data in enumerate(frame_list):\n",
    "        embedded += template.format(i, frame_format,\n",
    "                                    frame_data.replace('\\n', '\\\\\\n'))\n",
    "    return embedded\n",
    "\n",
    "\n",
    "@writers.register('html')\n",
    "class HTMLWriter(FileMovieWriter):\n",
    "    # we start the animation id count at a random number: this way, if two\n",
    "    # animations are meant to be included on one HTML page, there is a\n",
    "    # very small chance of conflict.\n",
    "    rng = random.Random()\n",
    "    exec_key = 'animation.ffmpeg_path'\n",
    "    args_key = 'animation.ffmpeg_args'\n",
    "    supported_formats = ['png', 'jpeg', 'tiff', 'svg']\n",
    "\n",
    "    @classmethod\n",
    "    def new_id(cls):\n",
    "        #return '%16x' % cls.rng.getrandbits(64)\n",
    "        return ''.join(cls.rng.choice(string.ascii_uppercase)\n",
    "                       for x in range(16))\n",
    "\n",
    "    def __init__(self, fps=30, codec=None, bitrate=None, extra_args=None,\n",
    "                 metadata=None, embed_frames=False, default_mode='loop'):\n",
    "        self.embed_frames = embed_frames\n",
    "        self.default_mode = default_mode.lower()\n",
    "\n",
    "        if self.default_mode not in ['loop', 'once', 'reflect']:\n",
    "            self.default_mode = 'loop'\n",
    "            warnings.warn(\"unrecognized default_mode: using 'loop'\")\n",
    "\n",
    "        self._saved_frames = list()\n",
    "        super(HTMLWriter, self).__init__(fps, codec, bitrate,\n",
    "                                         extra_args, metadata)\n",
    "\n",
    "    def setup(self, fig, outfile, dpi, frame_dir=None):\n",
    "        if os.path.splitext(outfile)[-1] not in ['.html', '.htm']:\n",
    "            raise ValueError(\"outfile must be *.htm or *.html\")\n",
    "\n",
    "        if not self.embed_frames:\n",
    "            if frame_dir is None:\n",
    "                frame_dir = outfile.rstrip('.html') + '_frames'\n",
    "            if not os.path.exists(frame_dir):\n",
    "                os.makedirs(frame_dir)\n",
    "            frame_prefix = os.path.join(frame_dir, 'frame')\n",
    "        else:\n",
    "            frame_prefix = None\n",
    "\n",
    "        super(HTMLWriter, self).setup(fig, outfile, dpi,\n",
    "                                      frame_prefix, clear_temp=False)\n",
    "\n",
    "    def grab_frame(self, **savefig_kwargs):\n",
    "        if self.embed_frames:\n",
    "            suffix = '.' + self.frame_format\n",
    "            f = InMemory()\n",
    "            self.fig.savefig(f, format=self.frame_format,\n",
    "                             dpi=self.dpi, **savefig_kwargs)\n",
    "            f.seek(0)\n",
    "            self._saved_frames.append(b64encode(f.read()).decode('ascii'))\n",
    "        else:\n",
    "            return super(HTMLWriter, self).grab_frame(**savefig_kwargs)\n",
    "\n",
    "    def _run(self):\n",
    "        # make a ducktyped subprocess standin\n",
    "        # this is called by the MovieWriter base class, but not used here.\n",
    "        class ProcessStandin(object):\n",
    "            returncode = 0\n",
    "            def communicate(self):\n",
    "                return ('', '')\n",
    "        self._proc = ProcessStandin()\n",
    "\n",
    "        # save the frames to an html file\n",
    "        if self.embed_frames:\n",
    "            fill_frames = _embedded_frames(self._saved_frames,\n",
    "                                           self.frame_format)\n",
    "        else:\n",
    "            # temp names is filled by FileMovieWriter\n",
    "            fill_frames = _included_frames(self._temp_names,\n",
    "                                           self.frame_format)\n",
    "\n",
    "        mode_dict = dict(once_checked='',\n",
    "                         loop_checked='',\n",
    "                         reflect_checked='')\n",
    "        mode_dict[self.default_mode + '_checked'] = 'checked'\n",
    "\n",
    "        interval = int(1000. / self.fps)\n",
    "\n",
    "        with open(self.outfile, 'w') as of:\n",
    "            of.write(JS_INCLUDE)\n",
    "            of.write(DISPLAY_TEMPLATE.format(id=self.new_id(),\n",
    "                                             Nframes=len(self._temp_names),\n",
    "                                             fill_frames=fill_frames,\n",
    "                                             interval=interval,\n",
    "                                             icons=_Icons(),\n",
    "                                             **mode_dict))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from matplotlib.animation import Animation\n",
    "import matplotlib.pyplot as plt\n",
    "import tempfile\n",
    "import random\n",
    "import os\n",
    "\n",
    "\n",
    "__all__ = ['anim_to_html', 'display_animation']\n",
    "\n",
    "\n",
    "class _NameOnlyTemporaryFile(object):\n",
    "    \"\"\"A context-managed temporary file which is not opened.\n",
    "    The file should be accessible by name on any system.\n",
    "    Parameters\n",
    "    ----------\n",
    "    suffix : string\n",
    "        The suffix of the temporary file (default = '')\n",
    "    prefix : string\n",
    "        The prefix of the temporary file (default = '_tmp_')\n",
    "    hash_length : string\n",
    "        The length of the random hash.  The size of the hash space will\n",
    "        be 16 ** hash_length (default=8)\n",
    "    seed : integer\n",
    "        the seed for the random number generator.  If not specified, the\n",
    "        system time will be used as a seed.\n",
    "    absolute : boolean\n",
    "        If true, return an absolute path to a temporary file in the current\n",
    "        working directory.\n",
    "    Example\n",
    "    -------\n",
    "    >>> with _NameOnlyTemporaryFile(seed=0, absolute=False) as f:\n",
    "    ...     print(f)\n",
    "    ...\n",
    "    _tmp_d82c07cd\n",
    "    >>> os.path.exists('_tmp_d82c07cd')  # file removed after context\n",
    "    False\n",
    "    \"\"\"\n",
    "    def __init__(self, prefix='_tmp_', suffix='', hash_length=8,\n",
    "                 seed=None, absolute=True):\n",
    "        rng = random.Random(seed)\n",
    "        self.name = '%s%0*x%s' % (prefix, hash_length,\n",
    "                                  rng.getrandbits(4 * hash_length), suffix)\n",
    "        if absolute:\n",
    "            self.name = os.path.abspath(self.name)\n",
    "\n",
    "    def __enter__(self):\n",
    "        return self\n",
    "\n",
    "    def __exit__(self, *exc_info):\n",
    "        if os.path.exists(self.name):\n",
    "            os.remove(self.name)\n",
    "\n",
    "\n",
    "def anim_to_html(anim, fps=None, embed_frames=True, default_mode='loop'):\n",
    "    \"\"\"Generate HTML representation of the animation\"\"\"\n",
    "    if fps is None and hasattr(anim, '_interval'):\n",
    "        # Convert interval in ms to frames per second\n",
    "        fps = 1000. / anim._interval\n",
    "\n",
    "    plt.close(anim._fig)\n",
    "    if hasattr(anim, \"_html_representation\"):\n",
    "        return anim._html_representation\n",
    "    else:\n",
    "        # tempfile can't be used here: we need a filename, and this\n",
    "        # fails on windows.  Instead, we use a custom filename generator\n",
    "        #with tempfile.NamedTemporaryFile(suffix='.html') as f:\n",
    "        with _NameOnlyTemporaryFile(suffix='.html') as f:\n",
    "            anim.save(f.name,  writer=HTMLWriter(fps=fps,\n",
    "                                                 embed_frames=embed_frames,\n",
    "                                                 default_mode=default_mode))\n",
    "            html = open(f.name).read()\n",
    "\n",
    "        anim._html_representation = html\n",
    "        return html\n",
    "\n",
    "\n",
    "def display_animation(anim, **kwargs):\n",
    "    \"\"\"Display the animation with an IPython HTML object\"\"\"\n",
    "    from IPython.display import HTML\n",
    "    return HTML(anim_to_html(anim, **kwargs))\n",
    "\n",
    "\n",
    "# This is the magic that makes animations display automatically in the\n",
    "# IPython notebook.  The _repr_html_ method is a special method recognized\n",
    "# by IPython.\n",
    "Animation._repr_html_ = anim_to_html"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
